import backtrader as bt
import backtrader.indicators as btind  # 导入策略分析模块
import pandas as pd
import datetime
import akshare as ak


def get_data(code):
    df = ak.stock_zh_a_daily(symbol=code, adjust='qfq')
    # df.date = pd.to_datetime(df.date, format='%Y-%m-%d', utc=True)
    # df.set_index('date', inplace=True)
    df.index = pd.to_datetime(df.date)
    df.sort_index(inplace=True)

    return df


# 自定义信号指标
class Mysignal(bt.Indicator):
    lines = ('signal',)  # 声明 signal 线，交易信号放在 signal line 上
    params = (('period', 30),)

    def __init__(self):
        self.lines.signal = self.data - bt.indicators.SMA(period=self.p.period)


codes = ['sh600276', 'sh600519', 'sh603288']
# 恒瑞医药
data = get_data(codes[0])

cerebro = bt.Cerebro()
st_date = datetime.datetime(2023, 3, 1)
end_date = datetime.datetime(2023, 5, 29)
datafeed = bt.feeds.PandasData(dataname=data, fromdate=st_date, todate=end_date)
cerebro.adddata(datafeed, name='600276.SH')

cerebro.add_signal(bt.SIGNAL_LONGSHORT, Mysignal, period=5)
cerebro.run()
cerebro.plot()
